# -*- coding: utf-8 -*-
import torch.nn as nn
import torch
from torchsummary import summary
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt
import os
from tqdm import tqdm

model_dict_path = "./model"
if not os.path.exists(model_dict_path):
    os.makedirs(model_dict_path)

batch_size = 16
train_path = './data/train'
test_path = './data/val'
num_workers = max(4, os.cpu_count() // 2)  # 使用50% CPU核心

# dataset_transform = transforms.Compose([transforms.ToTensor(),
#                                         transforms.Resize((227, 227))])
# 强化数据增强（AlexNet关键特性）
# 随机裁剪/翻转：增加数据多样性
# 颜色扰动：模拟光照变化，提升泛化性
# 标准化：加速收敛（使用ImageNet标准参数）

dataset_transform_train = transforms.Compose([
    transforms.RandomResizedCrop(227),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

dataset_transform_test = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(227),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])
flower_train = ImageFolder(train_path, transform=dataset_transform_train)
flower_test = ImageFolder(test_path, transform=dataset_transform_test)

train_loader = DataLoader(dataset=flower_train, batch_size=batch_size,
                          shuffle=True, num_workers=num_workers, pin_memory=True)
test_loader = DataLoader(dataset=flower_test, batch_size=batch_size,
                         shuffle=False, num_workers=num_workers, pin_memory=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Alexnet(nn.Module):
    def __init__(self, in_dim, num_classes):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=in_dim,
                      out_channels=96,
                      kernel_size=11,
                      stride=4,
                      padding=0),
            nn.ReLU(inplace=True),  # 正值保持原来结果，负值归零
            nn.MaxPool2d(3, 2),
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2)
        )
        self.fc = nn.Sequential(
            # nn.Linear(9216, 4096),
            nn.Linear(256*6*6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(4096, num_classes)
        )
        # 添加LRN层（AlexNet原始设计）
        self.conv.add_module("lrn1", nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75))

    def forward(self, x):
        # 输入数据送入卷积部分进行处理
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        output = self.fc(x)
        return output


model = Alexnet(in_dim=3, num_classes=5).to(device)
learning_rate = 1e-3
epochs = 3
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
loss_fn = nn.CrossEntropyLoss().to(device)


# 计算模型预测精度:测试集数据，模型
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 不进行梯度计算
    with torch.no_grad():
        # 模型是验证模式
        model.eval()
        # 获取每一个batch的数据，进行预测
        for images, labels in data_iter:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            # 获取预测结果
            _, predicts = torch.max(outputs.data, dim=-1)
            # 预测的次数
            total += labels.size(0)
            # 预测正确的个数
            correct += (predicts == labels).cpu().sum()
    # 获取准确率
    return correct / total


# 定义模型训练过程:指定数据集，优化器，损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=epochs):
    # 遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 每一个轮次的损失，预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 损失和的初始化
        train_loss_sum = 0
        # 迭代次数
        iter = 0
        for i, (images, labels) in enumerate(tqdm(data_loader, desc=f'Epoch {epoch + 1}/{epochs}', leave=True)):
            images = images.to(device)
            labels = labels.to(device)
            # 模型定义为训练模式
            model.train()
            # 对数据进行预测
            outputs = model(images)
            # 计算模型的损失
            loss = loss_fn(outputs, labels)
            # 在做反向传播前先清除网络状态
            optimizer.zero_grad()
            # 损失值进行反向传播
            loss.backward()
            # 参数迭代更新
            optimizer.step()
            # 求损失的和
            train_loss_sum += loss.item()
            # 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter += 1
        # 测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        print(
            'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
            .format(epoch, train_loss_sum / iter,
                    train_correct / train_accuracy_total, test_acc))
        torch.save(model.state_dict(), './model/alexnet.pth')
        print('保存成功')
    print('------------finish training-------------')


# 画图查看导入的图片

# for b, (images, labels) in enumerate(train_loader):
#     if b == 1:
#         fig, axes = plt.subplots(1, 2)
#         for i in range(batch_size):
#             axes[i].imshow(images[i].permute(1, 2, 0))
#             axes[i].set_title(labels[i].item())
#         plt.show()
#     elif b > 0:
#         break

if __name__ == '__main__':
    # net = Alexnet(3, 1000)
    # summary(model=net, input_size=(3, 227, 227), batch_size=1, device="cpu") #查看网络模型
    train()
